
// This file is part of the LITIV framework; visit the original repository at
// https://github.com/plstcharles/litiv for more information.
//
// Copyright 2015 Pierre-Luc St-Charles; pierre-luc.st-charles<at>polymtl.ca
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
/////////////////////////////////////////////////////////////////////////////
//
// VAP Trimodal Dataset Registration Tool
// =======================================
//
// This application registers RGB and Thermal images from the LITIV-VAP
// trimodal dataset using camera calibration parameters. The registered
// images are aligned with RGB as the reference frame.
//
// Features:
//  - Automatic lens distortion correction
//  - Stereo rectification (horizontal alignment)
//  - RGB/Thermal image registration
//  - Multiple output formats (separate + overlay)
//
// Usage:
//  vap_registration [options]
//
// Options:
//  --dataset <path>      Path to VAP dataset (default: ./datasets/litiv_vap/)
//  --output <path>       Output directory (default: ./datasets/litiv_vap/registered/)
//  --scene <name>        Scene name (default: "Scene 1")
//  --undistort           Enable lens distortion correction (default: true)
//  --rectify             Enable stereo rectification (default: false)
//  --no-depth            Skip depth channel loading
//  --display             Show real-time processing preview
//  --start <index>       Start frame index (default: 0)
//  --count <num>         Number of frames to process (default: all)
//
/////////////////////////////////////////////////////////////////////////////

#include "litiv/datasets.hpp"
#include <iostream>
#include <iomanip>

// Configuration parameters
struct RegistrationConfig {
    std::string sDatasetPath = "d:/Programming/ImageFusion/litiv/data/vap/rgbdt-stereo/";
    std::string sOutputPath = "d:/Programming/ImageFusion/litiv/data/vap/rgbdt-stereo/registered/";
    std::string sSceneName = "Scene 1";
    bool bUndistort = true;
    bool bHorizRectify = false;
    bool bLoadDepth = false;
    bool bDisplayPreview = false;
    size_t nStartIdx = 0;
    size_t nFrameCount = 0; // 0 = process all frames
};

// Parse command line arguments
bool parseArgs(int argc, char** argv, RegistrationConfig& config) {
    for(int i = 1; i < argc; ++i) {
        std::string arg = argv[i];

        if(arg == "--dataset" && i+1 < argc) {
            config.sDatasetPath = argv[++i];
            if(config.sDatasetPath.back() != '/' && config.sDatasetPath.back() != '\\')
                config.sDatasetPath += "/";
        }
        else if(arg == "--output" && i+1 < argc) {
            config.sOutputPath = argv[++i];
            if(config.sOutputPath.back() != '/' && config.sOutputPath.back() != '\\')
                config.sOutputPath += "/";
        }
        else if(arg == "--scene" && i+1 < argc) {
            config.sSceneName = argv[++i];
        }
        else if(arg == "--undistort") {
            config.bUndistort = true;
        }
        else if(arg == "--no-undistort") {
            config.bUndistort = false;
        }
        else if(arg == "--rectify") {
            config.bHorizRectify = true;
        }
        else if(arg == "--no-depth") {
            config.bLoadDepth = false;
        }
        else if(arg == "--display") {
            config.bDisplayPreview = true;
        }
        else if(arg == "--start" && i+1 < argc) {
            config.nStartIdx = std::stoul(argv[++i]);
        }
        else if(arg == "--count" && i+1 < argc) {
            config.nFrameCount = std::stoul(argv[++i]);
        }
        else if(arg == "--help" || arg == "-h") {
            return false;
        }
        else {
            std::cerr << "Unknown argument: " << arg << std::endl;
            return false;
        }
    }
    return true;
}

// Print usage information
void printUsage(const char* progName) {
    std::cout << "\n"
              << "VAP Trimodal Dataset Registration Tool\n"
              << "========================================\n\n"
              << "Usage: " << progName << " [options]\n\n"
              << "Options:\n"
              << "  --dataset <path>      Path to VAP dataset (default: ./data/vap/rgbdt-stereo/)\n"
              << "  --output <path>       Output directory (default: ./data/vap/rgbdt-stereo/registered/)\n"
              << "  --scene <name>        Scene name (default: \"Scene 1\")\n"
              << "  --undistort           Enable lens distortion correction (default: true)\n"
              << "  --no-undistort        Disable lens distortion correction\n"
              << "  --rectify             Enable stereo rectification (default: false)\n"
              << "  --no-depth            Skip depth channel loading\n"
              << "  --display             Show real-time processing preview\n"
              << "  --start <index>       Start frame index (default: 0)\n"
              << "  --count <num>         Number of frames to process (default: all)\n"
              << "  --help, -h            Show this help message\n\n"
              << "Examples:\n"
              << "  " << progName << " --scene \"Scene 1\" --undistort\n"
              << "  " << progName << " --dataset /path/to/vap/ --output /path/to/output/ --rectify\n"
              << "  " << progName << " --start 100 --count 50 --display\n\n";
}

int main(int argc, char** argv) {
    try {
        RegistrationConfig config;

        // Debug: Print all received arguments
        std::cout << "Received " << argc << " arguments:" << std::endl;
        for(int i = 0; i < argc; ++i) {
            std::cout << "  argv[" << i << "] = \"" << argv[i] << "\"" << std::endl;
        }
        std::cout << std::endl;

        // Parse command line arguments
        if(!parseArgs(argc, argv, config)) {
            printUsage(argv[0]);
            return argc > 1 ? -1 : 0;
        }

        std::cout << "\n========================================" << std::endl;
        std::cout << "VAP Dataset Registration Tool" << std::endl;
        std::cout << "========================================\n" << std::endl;

        std::cout << "Configuration:" << std::endl;
        std::cout << "  Dataset path:     " << config.sDatasetPath << std::endl;
        std::cout << "  Output path:      " << config.sOutputPath << std::endl;
        std::cout << "  Scene:            " << config.sSceneName << std::endl;
        std::cout << "  Undistort:        " << (config.bUndistort ? "YES" : "NO") << std::endl;
        std::cout << "  Rectify:          " << (config.bHorizRectify ? "YES" : "NO") << std::endl;
        std::cout << "  Load depth:       " << (config.bLoadDepth ? "YES" : "NO") << std::endl;
        std::cout << "  Display preview:  " << (config.bDisplayPreview ? "YES" : "NO") << std::endl;
        std::cout << std::endl;

        // Create dataset using VAP specialization
        // VAP is a VideoArray dataset, so it uses DatasetTask_Cosegm (cosegmentation)
        using DatasetType = lv::Dataset_<lv::DatasetTask_Cosegm,lv::Dataset_VAP_trimod2016,lv::NonParallel>;

        std::cout << "Initializing VAP dataset..." << std::endl;

        // Create dataset with VAP-specific parameters
        // VAP dataset uses a special create signature
        DatasetType::Ptr pDataset = DatasetType::create(
            "vap_registration_output",  // output directory name
            true,                        // save output
            false,                       // use evaluator
            config.bLoadDepth,          // load depth
            config.bUndistort,          // undistort
            config.bHorizRectify,       // horizontal rectify
            false,                       // eval disparities
            false,                       // load frame subset
            false,                       // eval only frame subset
            0,                          // eval temporal window size
            0,                          // load input masks
            1.0                         // scale factor
        );

        // Get the scene batch
        lv::IDataHandlerPtrArray vpBatches = pDataset->getBatches(false);
        lvAssert__(!vpBatches.empty(), "No batches found in dataset");

        // Find the requested scene by name
        DatasetType::WorkBatch* pBatch = nullptr;
        for(auto& pCurrBatch : vpBatches) {
            if(pCurrBatch->getName() == config.sSceneName) {
                pBatch = dynamic_cast<DatasetType::WorkBatch*>(pCurrBatch.get());
                break;
            }
        }

        if(!pBatch) {
            std::cerr << "Error: Could not find scene '" << config.sSceneName << "' in dataset" << std::endl;
            std::cerr << "Available scenes:" << std::endl;
            for(auto& pCurrBatch : vpBatches) {
                std::cerr << "  - " << pCurrBatch->getName() << std::endl;
            }
            return -1;
        }

        DatasetType::WorkBatch& oBatch = *pBatch;
        const size_t nTotalPackets = oBatch.getInputCount();

        std::cout << "Scene loaded: " << oBatch.getName() << std::endl;
        std::cout << "Total frames: " << nTotalPackets << std::endl;

        // Determine frame range to process
        size_t nStartIdx = config.nStartIdx;
        size_t nEndIdx = (config.nFrameCount > 0) ?
            std::min(nStartIdx + config.nFrameCount, nTotalPackets) : nTotalPackets;

        if(nStartIdx >= nTotalPackets) {
            std::cerr << "Error: Start index " << nStartIdx << " exceeds total frames " << nTotalPackets << std::endl;
            return -1;
        }

        std::cout << "Processing frames: " << nStartIdx << " to " << (nEndIdx-1) << std::endl;
        std::cout << std::endl;

        // Create output directories
        // Use the actual output path from the dataset
        std::string sActualOutputPath = oBatch.getOutputPath();
        std::string sSceneName = oBatch.getName();

        std::string sRGBOutDir = sActualOutputPath + "rgb/";
        std::string sThermalOutDir = sActualOutputPath + "thermal/";
        std::string sOverlayOutDir = sActualOutputPath + "overlay/";
        std::string sDepthOutDir = sActualOutputPath + "depth/";

        lv::createDirIfNotExist(sRGBOutDir);
        lv::createDirIfNotExist(sThermalOutDir);
        lv::createDirIfNotExist(sOverlayOutDir);
        if(config.bLoadDepth)
            lv::createDirIfNotExist(sDepthOutDir);

        std::cout << "Output directories created at: " << sActualOutputPath << std::endl;
        std::cout << "Starting registration..." << std::endl;
        std::cout << std::endl;

        // Processing loop
        oBatch.startProcessing();

        size_t nProcessedFrames = 0;
        auto tStart = std::chrono::high_resolution_clock::now();

        for(size_t nPacketIdx = nStartIdx; nPacketIdx < nEndIdx; ++nPacketIdx) {
            // Progress reporting
            if(nProcessedFrames % 10 == 0 || nProcessedFrames == 0) {
                auto tNow = std::chrono::high_resolution_clock::now();
                double dElapsed = std::chrono::duration<double>(tNow - tStart).count();
                double dFPS = (nProcessedFrames > 0) ? nProcessedFrames / dElapsed : 0.0;

                std::cout << "\rProcessing: "
                          << std::setw(5) << nPacketIdx
                          << " / " << (nEndIdx-1)
                          << "  [" << std::fixed << std::setprecision(2) << dFPS << " fps]"
                          << std::flush;
            }

            // Get registered input images (already undistorted and rectified by the dataset parser)
            // For VideoArray datasets (Cosegm task), use getInputArray
            std::vector<cv::Mat> vInputs = oBatch.getInputArray(nPacketIdx);

            if(vInputs.empty()) {
                std::cerr << "\nWarning: No data for packet " << nPacketIdx << std::endl;
                continue;
            }

            // Input streams:
            // vInputs[0] = RGB (CV_8UC3, undistorted)
            // vInputs[1] = Thermal (CV_8UC1, undistorted & aligned)
            // vInputs[2] = Depth (if loaded)

            cv::Mat oRGB = vInputs[0];
            cv::Mat oThermal = (vInputs.size() > 1) ? vInputs[1] : cv::Mat();
            cv::Mat oDepth = (vInputs.size() > 2) ? vInputs[2] : cv::Mat();

            // When rectification is enabled, the dataset rotates images 90 degrees CW
            // Original: transpose + flip(1) = rotate 90 CW
            // To undo (rotate 90 CCW): flip(0) + transpose
            if(config.bHorizRectify) {
                if(!oRGB.empty()) {
                    cv::flip(oRGB, oRGB, 1);           // Vertical flip
                    cv::transpose(oRGB, oRGB);         // Transpose
                }
                if(!oThermal.empty()) {
                    cv::flip(oThermal, oThermal, 1);
                    cv::transpose(oThermal, oThermal);
                }
                if(!oDepth.empty()) {
                    cv::flip(oDepth, oDepth, 1);
                    cv::transpose(oDepth, oDepth);
                }
            }

            // Generate output filenames
            std::string sFrameName = cv::format("%05d.png", (int)nPacketIdx);

            // Save registered RGB
            if(!oRGB.empty()) {
                cv::imwrite(sRGBOutDir + sFrameName, oRGB);
            }

            // Save registered Thermal
            if(!oThermal.empty()) {
                cv::imwrite(sThermalOutDir + sFrameName, oThermal);
            }

            // Save depth if loaded
            if(config.bLoadDepth && !oDepth.empty()) {
                cv::imwrite(sDepthOutDir + sFrameName, oDepth);
            }

            // Create overlay visualization (RGB + pseudocolored thermal)
            if(!oRGB.empty() && !oThermal.empty()) {
                cv::Mat oThermalColor;
                cv::applyColorMap(oThermal, oThermalColor, cv::COLORMAP_JET);

                cv::Mat oOverlay;
                cv::addWeighted(oRGB, 0.6, oThermalColor, 0.4, 0, oOverlay);

                // Add text label
                std::string sLabel = cv::format("Frame %05d | RGB-Thermal Registration", (int)nPacketIdx);
                cv::putText(oOverlay, sLabel, cv::Point(10, 25),
                           cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(255, 255, 255), 2);
                cv::putText(oOverlay, sLabel, cv::Point(10, 25),
                           cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(0, 0, 0), 1);

                cv::imwrite(sOverlayOutDir + sFrameName, oOverlay);

                // Display preview if enabled
                if(config.bDisplayPreview) {
                    cv::imshow("RGB", oRGB);
                    cv::imshow("Thermal", oThermal);
                    cv::imshow("Overlay", oOverlay);

                    int nKey = cv::waitKey(1);
                    if(nKey == 27 || nKey == 'q') { // ESC or 'q' to quit
                        std::cout << "\n\nUser interrupted. Stopping..." << std::endl;
                        break;
                    }
                }
            }

            ++nProcessedFrames;
        }

        oBatch.stopProcessing();

        // Final statistics
        auto tEnd = std::chrono::high_resolution_clock::now();
        double dTotalTime = std::chrono::duration<double>(tEnd - tStart).count();
        double dAvgFPS = nProcessedFrames / dTotalTime;

        std::cout << "\n\n========================================" << std::endl;
        std::cout << "Registration Complete!" << std::endl;
        std::cout << "========================================" << std::endl;
        std::cout << "Processed frames:  " << nProcessedFrames << std::endl;
        std::cout << "Total time:        " << std::fixed << std::setprecision(2) << dTotalTime << " seconds" << std::endl;
        std::cout << "Average speed:     " << std::fixed << std::setprecision(2) << dAvgFPS << " fps" << std::endl;
        std::cout << "Output location:   " << config.sOutputPath << config.sSceneName << "/" << std::endl;
        std::cout << std::endl;

        std::cout << "Output structure:" << std::endl;
        std::cout << "  rgb/       - Registered RGB images (reference frame)" << std::endl;
        std::cout << "  thermal/   - Registered thermal images (aligned to RGB)" << std::endl;
        std::cout << "  overlay/   - RGB + thermal overlay visualization" << std::endl;
        if(config.bLoadDepth)
            std::cout << "  depth/     - Depth images" << std::endl;
        std::cout << std::endl;

        if(config.bDisplayPreview)
            cv::destroyAllWindows();

    }
    catch(const lv::Exception& e) {
        std::cerr << "\n\nLITIV Exception: " << e.what() << std::endl;
        return -1;
    }
    catch(const cv::Exception& e) {
        std::cerr << "\n\nOpenCV Exception: " << e.what() << std::endl;
        return -1;
    }
    catch(const std::exception& e) {
        std::cerr << "\n\nStandard Exception: " << e.what() << std::endl;
        return -1;
    }
    catch(...) {
        std::cerr << "\n\nUnknown exception caught!" << std::endl;
        return -1;
    }

    return 0;
}
